import time
import numpy as np

policy_path = f'DP_policy_2v1_small/train/'
adj_path = f'adj_file/adj_file_train_small/'

for index in range(76):
    starttime = time.time()
    opponent_file = adj_path + str(index) + '.txt'
    ajacent_file = adj_path + 'adjacent_matrix_img_' + str(index).zfill(3) + '.png.txt'
    # ajacent_file = adj_path + 'adjacent_matrix_img_' + str(index) + '.png.txt'

    # 读取npy文件
    # opponent_policy = np.load('DP_policy/opponent_policy_'+str(index)+'.npy')
    # pursuer_policy = np.load('DP_policy/pursuer_policy_'+str(index)+'.npy')
    # Global_config.opponent_policy.append(opponent_policy)
    # Global_config.pursuer_policy.append(pursuer_policy)

    # 存储成npy文件
    data_list = []
    with open(opponent_file, 'r') as file:
        for line in file:
            row = [int(x) for x in line.strip().split()]
            data_list.append(row)
    # opponent_policy = np.array(data_list)

    data_array = np.array(data_list)
    column_1 = data_array[:, 0]
    column_2 = data_array[:, 1]
    column_3 = data_array[:, 2]
    # column_4 = data_array[:, 3]

    with open(ajacent_file, "r") as file:
        first_line = file.readline().strip()
        elements = first_line.split()
        num = int(elements[0])
    # opponent_policy = opponent_policy.reshape(num, num, num, num)

    # 3v1
    # opponent_policy = column_1.reshape(num, num, num, num)
    # pursuer_1_policy = column_2.reshape(num, num, num, num)
    # pursuer_2_policy = column_3.reshape(num, num, num, num)
    # pursuer_3_policy = column_4.reshape(num, num, num, num)

    # 2v1
    opponent_policy = column_1.reshape(num, num, num)
    pursuer_1_policy = column_2.reshape(num, num, num)
    pursuer_2_policy = column_3.reshape(num, num, num)

    pursuer_policy = []
    pursuer_policy.append(pursuer_1_policy)
    pursuer_policy.append(pursuer_2_policy)
    # pursuer_policy.append(pursuer_3_policy)

    pursuer_policy = np.array(pursuer_policy)

    # np.save('DP_policy/opponent_policy_'+str(index)+'.npy', opponent_policy)
    np.save(policy_path + 'opponent_policy_'+str(index)+'.npy', opponent_policy)
    np.save(policy_path + 'pursuer_policy_'+str(index)+'.npy', pursuer_policy)
    print('存储成功')
    print("opponent data shape:", opponent_policy.shape)
    print("pursuer data shape:", pursuer_policy.shape)
    
    print("Finish", index, "time cost", time.time()-starttime)
## ---------------------------------------------------------------
# # 编号0->000
# import os

# # 文件夹路径
# folder_path = "DP_policy_2v1_small/train"

# # 重命名文件
# for i in range(76):
#     old_name = os.path.join(folder_path, f"opponent_policy_{i}.npy")
#     new_name = os.path.join(folder_path, f"opponent_policy_{i:03d}.npy")
#     if os.path.exists(old_name):
#         os.rename(old_name, new_name)
#------------------------------------------------------
# # 重新编号
# import os

# # 文件夹路径
# folder_path = "DP_policy_2v1_small/train"

# # 获取文件列表并排序
# files = [f for f in os.listdir(folder_path) if f.startswith('opponent')]
# files = sorted(files)

# # 新文件名模板
# new_names = []

# # 穿插两个列表生成新顺序
# for i in range(76):
#     new_names.append((files[i], f"opponent_policy_{2 * i + 1}.npy"))
#     # new_names.append((files[i], f"adjacent_matrix_img_{2 * i:03d}.png.txt"))
# new_names.reverse()
# print(new_names)
# # 重命名文件
# for old_name, new_name in new_names:
#     old_path = os.path.join(folder_path, old_name)
#     new_path = os.path.join(folder_path, new_name)
#     os.rename(old_path, new_path)
